import time
start =time.perf_counter()

import gym
from RL_Brain import PolicyGradient
import matplotlib.pyplot as plt
t=10                
DISPLAY_REWARD_THRESHOLD =68000

RENDER = False  
# rendering wastes time 
# 当回合总 reward 大于 68000 时显示模拟窗口
#env = gym.make('CartPole-v0')
env = gym.make('GridWorld-v0')  # CartPole 这个模拟
env.seed(1)     
env = env.unwrapped    # 取消限制

print(env.action_space)    # 显示可用 action
print(env.observation_space)    # 显示可用 state 的 observation
print(env.observation_space.high)   # 显示 observation 最高值
print(env.observation_space.low)    # 显示 observation 最低值

# 定义角色
RL = PolicyGradient(
    #n_actions=1,
    n_actions=env.action_space.shape[0],
    n_features=env.observation_space.shape[0],
    #n_features=3,
    learning_rate=0.1,
    reward_decay=0.89,
    
)

for i_episode in range(300):

    observation = env.reset()

    while True:
        if RENDER: env.render()    # 刷新环境

        action = RL.choose_action(observation)     # 选行为

        observation_, reward, done, info = env.step(action)     # 获取下一个 state

        RL.store_transition(observation, action, reward)

        if done:
            ep_rs_sum = sum(RL.ep_rs)

            if 'running_reward' not in globals():
                running_reward = ep_rs_sum
            else:
                running_reward = running_reward * 0.99 + ep_rs_sum * 0.01
            if running_reward > DISPLAY_REWARD_THRESHOLD: RENDER = True     # 判断是否显示模拟
            print("episode:", i_episode, "  reward:", int(running_reward))
            if i_episode==t: RENDER = True
            vt = RL.learn()

            if i_episode == 0:
                plt.plot(vt)    # plot 这个回合的 vt
                plt.xlabel('episode steps')
                plt.ylabel('normalized state-action value')
                plt.show()
            
            break

        observation = observation_



end = time.perf_counter()
print("程序运行时间为:")
print (end -start)
#print("运行成功次数为:")

